import matplotlib
matplotlib.use('QtAgg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from matplotlib import animation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from matplotlib.colors import to_hex
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
from PyQt5.QtWidgets import QMainWindow, QWidget, QVBoxLayout
from PyQt5.QtCore import Qt

Colors = [
    '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f',
    '#bcbd22', '#17becf', '#aec7e8', '#ffbb78', '#98df8a', '#ff9896', '#c5b0d5', '#c49c94',
    '#f7b6d2', '#c7c7c7', '#dbdb8d', '#9edae5', '#393b79', '#637939', '#8c6d31', '#843c39',
    '#ad494a', '#8c663b', '#8c6279', '#be6dd6', '#de9ed6'
]

class Animation3D(QMainWindow):
    def __init__(self, maps, starts, goals, paths):
        super().__init__()
        self.setWindowTitle('3D Path Animation')
        
        # 创建主窗口部件和布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout(main_widget)
        
        # 创建图形和画布
        self.fig = plt.figure(figsize=(12, 8))
        canvas = FigureCanvasQTAgg(self.fig)
        layout.addWidget(canvas)
        
        # 设置窗口大小和位置
        self.resize(1200, 800)
        self.setWindowState(Qt.WindowMaximized)
        
        # 初始化3D图表
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.ax.set_box_aspect([0.5, 0.5, 0.5])
        self.ax.set_zlim3d(0, len(maps) + 1)

        """
        修改坐标转换逻辑
        :param maps: 多层地图列表
        :param starts: 多层起点列表
        :param goals: 多层终点列表
        :param paths: 多层路径列表
        """
        # 修正坐标转换，不再翻转和转置矩阵
        self.maps = maps
        # 保持原始坐标顺序
        self.starts = starts
        self.goals = goals
        self.paths = paths
        self.layers = len(self.maps)

        # 计算地图比例
        aspect = len(self.maps[0]) / len(self.maps[0][0])

        # 3D 绘图
        self.ax.set_box_aspect([0.5, 0.5, 0.5])  # 试着调小 Z 轴的比例
        self.ax.set_zlim3d(0, self.layers + 1)  # 设定 z 轴范围

        self.T = 0
        self.agents = {}
        self.agent_names = {}
        self.artists = []

        # 生成层级颜色
        colormap = plt.cm.get_cmap("tab10", self.layers)
        layer_colors = [colormap(i) for i in range(self.layers)]

        # 遍历所有层
        for layer in range(self.layers):
            # 绘制每个路径的轨迹
            for i, path in enumerate(self.paths[layer]):
                path_points = np.array(path)
                # 绘制路径轨迹，使用半透明的颜色
                self.ax.plot(path_points[:, 1], path_points[:, 0], [layer + 1] * len(path_points),
                           color=layer_colors[layer], alpha=0.3, linestyle='--', linewidth=1.5)

            for x in range(len(self.maps[layer])):
                for y in range(len(self.maps[layer][0])):
                    if self.maps[layer][x][y] == 1:  # 如果是障碍物
                        # 立方体的 8 个顶点，保持原始坐标系
                        v = [
                            [y - 0.5, x - 0.5, layer],  # 底面四个点
                            [y + 0.5, x - 0.5, layer],
                            [y + 0.5, x + 0.5, layer],
                            [y - 0.5, x + 0.5, layer],
                            [y - 0.5, x - 0.5, layer + 1],  # 顶面四个点
                            [y + 0.5, x - 0.5, layer + 1],
                            [y + 0.5, x + 0.5, layer + 1],
                            [y - 0.5, x + 0.5, layer + 1]
                        ]

                        # 立方体的 6 个面，每个面由 4 个点构成
                        faces = [
                            [v[0], v[1], v[2], v[3]],  # 底面
                            [v[4], v[5], v[6], v[7]],  # 顶面
                            [v[0], v[1], v[5], v[4]],  # 前面
                            [v[2], v[3], v[7], v[6]],  # 后面
                            [v[1], v[2], v[6], v[5]],  # 右侧
                            [v[0], v[3], v[7], v[4]]  # 左侧
                        ]

                        # 创建 3D 立方体
                        poly = Poly3DCollection(faces, facecolor='black', edgecolor='black', alpha=0.3)
                        self.ax.add_collection3d(poly)
            layer_color = layer_colors[layer]  # 当前层颜色

            # 画终点
            for goal in self.goals[layer]:
                x, y = goal
                self.ax.scatter(y, x, layer + 1, color=layer_color, marker='*', s=150, label='Goal')

            # 画路径和起点
            for i in range(len(self.paths[layer])):
                name = f"L{layer}_A{i}"
                start_x, start_y = self.starts[layer][i]

                self.agents[(layer, i)] = self.ax.scatter(start_y, start_x, layer + 1,
                                                          color=layer_color, s=100)
                self.agent_names[(layer, i)] = self.ax.text(start_y, start_x + 0.25,
                                                            layer + 1, name, horizontalalignment='center')

                # 计算最大时间步
                self.T = max(self.T, len(self.paths[layer][i]) - 1)
                self.artists.append(self.agent_names[(layer, i)])

        def set_axes_equal(ax):
            '''Set 3D plot axes to equal scale.'''
            x_limits = ax.get_xlim3d()
            y_limits = ax.get_ylim3d()
            z_limits = ax.get_zlim3d()

            x_range = np.abs(x_limits[1] - x_limits[0])
            y_range = np.abs(y_limits[1] - y_limits[0])
            z_range = np.abs(z_limits[1] - z_limits[0])

            range_max = max(x_range, y_range, z_range)

            ax.set_xlim3d([x_limits[0], x_limits[0] + range_max])
            ax.set_ylim3d([y_limits[0], y_limits[0] + range_max])
            ax.set_zlim3d([z_limits[0], z_limits[0] + range_max])

        set_axes_equal(self.ax)

        self.ax.view_init(elev=30, azim=45)  # 设置初始视角
        self.ax.set_xlabel('X')
        self.ax.set_ylabel('Y')
        self.ax.set_zlabel('Layer')

        self.animation = animation.FuncAnimation(self.fig, self.animate_func,
                                                 init_func=self.init_func,
                                                 frames=int(self.T + 1) * 10,
                                                 interval=100,
                                                 blit=False)

        text_lines = []
        for layer_idx in range(self.layers):
            num_agents = len(self.starts[layer_idx])
            layer_color = to_hex(layer_colors[layer_idx])
            text_lines.append(
                f"Layer {layer_idx}: {num_agents} agents | Color: {layer_color}")

        # 添加半透明文本背景框
        self.fig.text(0.05, 0.75,  # 调整位置
                      "\n".join(text_lines),
                      fontsize=9,
                      linespacing=1.5,
                      bbox=dict(boxstyle="round",
                                facecolor="white",
                                alpha=0.7,
                                edgecolor="lightgray"),
                      zorder=10)

        # 保持动画引用
        self._anim = self.animation

        # 添加窗口关闭事件处理
        self.destroyed.connect(self.cleanup)
        
    def closeEvent(self, event):
        """窗口关闭时的清理"""
        self.cleanup()
        super().closeEvent(event)
        
    def cleanup(self):
        """清理资源"""
        if hasattr(self, 'animation') and self.animation:
            self.animation.event_source.stop()
            self.animation = None
        
        if hasattr(self, 'fig'):
            plt.close(self.fig)
            self.fig = None

    def init_func(self):
        for (layer, i), agent in self.agents.items():
            init_pos = self.paths[layer][i][0]
            agent.set_offsets([[init_pos[0], init_pos[1]]])
            agent.set_3d_properties(zs=[layer + 1], zdir='z')
            self.agent_names[(layer, i)].set_position((init_pos[0], init_pos[1] + 0.5))
            self.agent_names[(layer, i)].set_3d_properties(layer + 1, zdir='z')
        return self.artists

    def animate_func(self, t):
        for (layer, i), agent in self.agents.items():
            pos = self.get_state(t / 10, self.paths[layer][i])
            # 交换x,y坐标显示
            agent.set_offsets([[pos[1], pos[0]]])
            agent.set_3d_properties(zs=[layer + 1], zdir='z')
            self.agent_names[(layer, i)].set_position((pos[1], pos[0] + 0.5))
            self.agent_names[(layer, i)].set_3d_properties(layer + 1, zdir='z')
            
        return self.artists

    @staticmethod
    def get_state(t, path):
        if int(t) <= 0:
            return np.array(path[0])
        elif int(t) >= len(path):
            return np.array(path[-1])
        else:
            pos_last = np.array(path[int(t) - 1])
            pos_next = np.array(path[int(t)])
            pos = (pos_next - pos_last) * (t - int(t)) + pos_last
            return pos

    def save(self, file_name, speed=1):
        self.animation.save(file_name, fps=10 * speed, dpi=200, savefig_kwargs={"pad_inches": 0, "bbox_inches": "tight"})

    def show(self):
        """重写show方法"""
        super().show()
        # 确保窗口被激活并置于前台
        self.activateWindow()
        self.raise_()
